from PySide6.QtWidgets import QMainWindow, QPushButton, QVBoxLayout, QWidget, QFileDialog,QApplication
from PySide6.QtCore import QThread, Signal
import sys
from hats_detection.video_worker import process_video
from hats_detection.image_worker import process_image

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Hats Detection Application")

        # 创建按钮
        self.process_video_button = QPushButton("Process Video", self)
        self.process_video_button.clicked.connect(self.process_video)

        self.process_image_button = QPushButton("Process Image", self)
        self.process_image_button.clicked.connect(self.process_image)

        # 创建布局
        layout = QVBoxLayout()
        layout.addWidget(self.process_video_button)
        layout.addWidget(self.process_image_button)

        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

    def process_video(self):
        video_path, _ = QFileDialog.getOpenFileName(self, "Select Video", "", "Video Files (*.mp4 *.avi)")
        if video_path:
            process_video(video_path)

    def process_image(self):
        image_path, _ = QFileDialog.getOpenFileName(self, "Select Image", "", "Image Files (*.jpg *.png *.bmp)")
        if image_path:
            process_image(image_path)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())